
from hashlib import sha1
from xmlrpclib import ServerProxy, Binary

#URL = 'http://localhost:8080/'
URL = 'http://openlookup.appspot.com/'
#URL = 'http://opendht.nyuld.net:5851/'

#
# Client API
#

p = ServerProxy(URL)

class Error(Exception): pass

def get(key, maxvals=10, placemark='', application='test'):
    values, placemark = p.get(Binary(key), maxvals, Binary(placemark), application)
    values = [v.data for v in values]
    placemark = placemark.data
    assert len(values) <= maxvals
    return values, placemark

def get_details(key, maxvals=10, placemark='', application='test'):
    results, placemark = p.get_details(Binary(key), maxvals, Binary(placemark), application)
    results = [(value.data, ttl, hash_type, secret_hash.data)
               for (value, ttl, hash_type, secret_hash) in results]
    placemark = placemark.data
    return results, placemark

def put(key, value, secret='', ttl=10, application='test'):
    if not secret:
        status = p.put(Binary(key), Binary(value), ttl, application)
    else:
        secret_hash = sha1(secret).digest()
        status = p.put_removable(Binary(key), Binary(value),
                                 'SHA', Binary(secret_hash), ttl, application)
    if status: raise Error

def rm(key, value, secret, ttl=10, application='test'):
    value_hash = sha1(value).digest()
    status = p.rm(Binary(key), Binary(value_hash),
                  'SHA', Binary(secret), ttl, application)
    if status: raise Error

#
# Helpers
#

def get_all(key, batch_size=10):
    values = []
    placemark = ''
    while True:
        batch, placemark = get(key, batch_size, placemark)
        values.extend(batch)
        if not placemark:
            return sorted(values)

#
# Tests
#

def test_get():
    print 'test_get()'
    put('shape', 'round')
    put('shape', 'square')
    put('shape', 'square')
    put('shape', 'square', 'xyzzy')
    put('shape', 'square', 'xyzzy')
    put('shape', 'square', 'plugh')
    put('size', 'medium')
    expected = ['round', 'square', 'square', 'square']
    for batch_size in range(1, 6):
        assert get_all('shape', batch_size) == expected

def test_get_details():
    print 'test_get_details()'
    put('one', '1')
    results, placemark = get_details('one')
    value, ttl, hash_type, secret_hash = results[0]
    assert value == '1'
    assert 0 < ttl <= 10
    assert hash_type == ''
    assert secret_hash == ''

    put('two', '2', 'secret')
    results, placemark = get_details('two')
    value, ttl, hash_type, secret_hash = results[0]
    assert value == '2'
    assert 0 < ttl <= 10
    assert hash_type == 'SHA'
    assert secret_hash == sha1('secret').digest()

def test_put():
    print 'test_put()'
    put('answer', '42')
    put('answer', '42', ttl=20)
    assert get_all('answer') == ['42']
    results, placemark = get_details('answer')
    value, ttl, _, _ = results[0]
    assert ttl > 10

def test_rm():
    print 'test_rm()'
    k = sha1('color').digest()
    v = sha1('green').digest()
    s = sha1('sesame').digest()

    put(k, v, s)
    assert get_all(k) == [v]

    rm(k, v, s)
    assert get_all(k) == []

if __name__ == '__main__':
    test_get()
    test_get_details()
    test_put()
    test_rm()
